import torch 
from functools import partial
from torch.utils.data import DataLoader
from models.Update import DatasetSplit
from utils.options import args_parser
from utils.get_dataset import get_dataset
from sim.cka_torch import *
# from models.ResNet_Dery import *
from sim.block_meta import *
from models.resnet import *

def get_rep(args,net_list,dataset_global=None,dict_global = None):
    if args.num_groups != len(MODEL_ZOO):
        exit("error")
    container = dict.fromkeys(MODEL_ZOO)
    for key in container.keys():
        container[key] = dict.fromkeys(MODEL_BLOCKS[key])
        
    for key1 in container.keys():
        for key2 in container[key1].keys():
            container[key1][key2] = []

    train_data = DataLoader(DatasetSplit(dataset_global,dict_global),args.server_bs,shuffle=True)
    with torch.no_grad():

        def forward_hook(module,input,output,model_key,block_key):
            container[model_key][block_key].append(output)

        # for id in range(args.num_groups):
        #     hook = net_list[id]
        for id,key1 in enumerate(MODEL_ZOO) :
            for id,key2 in enumerate (MODEL_BLOCKS[key1]):
                hook = net_list[id]
                layers = key2.split('.')
                for lay in layers:
                    hook = getattr(hook,lay)
                hook.register_forward_hook(partial(forward_hook,model_key = key1),block_key = MODEL_BLOCKS[key1][key2])
            net_list[id].eval()
            for images in enumerate(train_data):
                images = images.to(args.device)
                net_list[id](images)

    return container

def compute_sim(args,container,model1,model2):
    data1 = container[model1]
    data2 = container[model2]
    sim = {}
    for block1 in enumerate(data1):
        for block2 in enumerate(data2):
            sim = similarity_pair_batch_cka(block1,block2,bs=args.bs)
     


if __name__ == '__main__': 
    args = args_parser()
    dataset_train,dataset_test,dict_users ,dict_global= get_dataset(args)
    net_list = []
    net_list.append(ResNet18_cifar10(num_classes = args.num_classes))
    net_list.append(ResNet18_cifar10(num_classes = args.num_classes))
    container = None
    container = get_rep(args=args,net_list=net_list,dataset_global=dataset_train,dict_global=dict_global)
    print(container)
    
